# Copyright (C) 2025 AIDC-AI
# This project is licensed under the MIT License (SPDX-License-identifier: MIT).

import json
import keyword
import os
from pathlib import Path
import re
import shutil
from pydantic import Field
from core import mcp, logger
from manager.workflow_manager import CUSTOM_WORKFLOW_DIR, workflow_manager
from utils.file_util import download_files


@mcp.tool(name="save_workflow_tool")
async def save_workflow_tool(
    workflow_url: str = Field(description="The workflow to save, must be a url"),
    uploaded_filename: str = Field(description="use file attachment name that user uploaded or a name specified by user, must be english and no file suffix"),
):
    """
    Add or update a workflow to MCP tools.
    """
    def error(msg: str):
        return json.dumps({ "success": False, "error": msg })
    
    try:
        # 合法格式：字母或下划线开头，其后可包含字母、数字、下划线
        pattern = r'^[A-Za-z_][A-Za-z0-9_]*$'
        
        if not re.match(pattern, uploaded_filename):
            return error(
                f"上传文件的名称格式不符合要求：仅允许使用字母、数字、下划线，且必须以字母或下划线开头。"
            )
        
        if keyword.iskeyword(uploaded_filename):
            return error(
                f"上传文件的名称不能为 Python 关键字：'{uploaded_filename}'。"
            )

        with download_files(workflow_url) as temp_workflow_path:
            return workflow_manager.load_workflow(temp_workflow_path, tool_name=uploaded_filename)
            
    except Exception as e:
        logger.error(f"保存工作流失败: {e}")
        return error(f"保存工作流失败: {str(e)}")
        
@mcp.tool(name="reload_workflows_tool")
async def reload_workflows_tool():
    """
    Reload all mcp tools that generated by workflows.
    """
    return workflow_manager.reload_all_workflows()
        
@mcp.tool(name="list_workflows_tool")
async def list_workflows_tool():
    """
    List all mcp tools that generated by workflows.
    
    Call this tool when user want to list all mcp tools, such as:
    - query all tools
    - list all mcp tools
    - how many tools are there
    """
    workflow_names = list(workflow_manager.loaded_workflows.keys())
    workflow_names.sort()
    return workflow_names

@mcp.tool(name="remove_workflow_tool")
async def remove_workflow_tool(
    workflow_name: str = Field(description="The name of the workflow to remove"),
):
    """
    Remove a mcp tool that generated by a workflow.
    """
    return workflow_manager.unload_workflow(workflow_name)
